package org.fujay.commons.langchain4j.rag.augmentor;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.filter.Filter;
import lombok.Builder;
import lombok.Data;
import lombok.experimental.Accessors;

import java.util.List;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;

/**
 * 增强检索接口
 * 提供混合检索、缓存管理等功能
 */
public interface EnhancedRetrieval extends ContentRetriever {

    /**
     * 执行检索
     *
     * @param query 查询对象
     * @return 检索结果列表
     */
    List<Content> retrieve(Query query);

    /**
     * 预热缓存
     *
     * @param queries 预热查询列表
     */
    void warmupCache(List<Query> queries);

    /**
     * 清除缓存
     */
    void clearCache();

    /**
     * 获取缓存统计信息
     *
     * @return 缓存统计信息
     */
    CacheStats getCacheStats();

    /**
     * 缓存统计信息
     */
    @Data
    @Builder
    @Accessors(chain = true)
    class CacheStats {
        /**
         * 缓存命中次数
         */
        private long hitCount;
        /**
         * 缓存未命中次数
         */
        private long missCount;
        /**
         * 缓存驱逐次数
         */
        private long evictionCount;
        /**
         * 当前缓存中的条目数量
         */
        private long size;
    }

    /**
     * 检索配置
     */
    @Data
    @Builder
    @Accessors(chain = true)
    class Config {
        /**
         * 检索器显示名称，用于日志和监控
         */
        private String displayName;

        /**
         * 向量模型，用于生成文本的向量表示
         */
        private EmbeddingModel embeddingModel;

        /**
         * 向量存储，用于存储和检索向量
         */
        private EmbeddingStore<TextSegment> embeddingStore;

        /**
         * 获取最大结果数量的函数
         */
        private Function<Query, Integer> maxResultsProvider;

        /**
         * 获取最小相似度分数的函数
         */
        private Function<Query, Double> minScoreProvider;

        /**
         * 获取过滤条件的函数
         */
        private Function<Query, Filter> filterProvider;

        /**
         * 是否启用混合检索（同时使用语义检索和关键词检索）
         */
        private boolean enableHybridSearch;

        /**
         * 是否启用BM25算法
         */
        private boolean enableBM25;

        /**
         * 语义检索的权重
         */
        private double semanticWeight;

        /**
         * 关键词检索的权重
         */
        private double keywordWeight;

        /**
         * 缓存大小
         */
        private int cacheSize;

        /**
         * 缓存过期时间（秒）
         */
        private long cacheExpireTime;

        /**
         * 创建默认配置
         *
         * @return 默认配置实例
         */
        public static Config defaultConfig(EmbeddingModel embeddingModel, EmbeddingStore<TextSegment> embeddingStore) {
            return Config.builder()
                    .displayName("默认检索器")
                    .embeddingModel(embeddingModel)
                    .embeddingStore(embeddingStore)
                    .maxResultsProvider(query -> 10)
                    .minScoreProvider(query -> 0.7)
                    .filterProvider(query -> null)
                    .enableHybridSearch(true)
                    .enableBM25(true)
                    .semanticWeight(0.7)
                    .keywordWeight(0.3)
                    .cacheSize(1000)
                    .cacheExpireTime(TimeUnit.HOURS.toSeconds(1))
                    .build();
        }
    }
}
